'''
Created on Feb 14, 2012

@author: oabalbin
takes a example analysis file.

It takes a file with a list of exome bam files
for tumor/normal samples and 
creates an analysis file for each pair.

TCGA-EJ-5508-10A-01D-1577-08.3.chr21.bam
./TCGA-CH-5768/unknown_sample_type/TCGA-CH-5768-11A-01D-1576-08/SRZ012613/provisional/C529.TCGA-CH-5768-11A-01D-1576-08.6.bam

'''
import os
import sys
import subprocess
import numpy as np
#import xml.etree.cElementTree as etree
import lxml.etree as etree
#import xml.etree.cElementTree as etree
from optparse import OptionParser
from collections import defaultdict


class tcgaRecord:
    def  __init__(self,fields,attributes):
        self.myatt=defaultdict()
        self.att=[]
        for f,a in zip(fields,attributes):
            self.myatt[f]=a
            self.att.append(f)
            
        self.submitted_sample_id=self.myatt["patiendid"]
        self.subjectid=self.submitted_sample_id.split('-')[:2]
        
    def print_myatt(self):
        return self.att
    def print_myself(self):
        l=[]
        for f in self.att:
            l.append( self.myatt[f] )
        return l
    def is_tumor(self):
        if self.myatt['is_sample_tumor']!='normal':
            return True
        else:
            return False
    def sample_name(self):
        return self.myatt['submitted_sample_id']
    def center(self):
        return self.myatt['center']
    def disease(self):
        return self.myatt['disease']


def read_bam_file_list(file,rootpath):
    '''
    returns a patient dictionary with the files for the tumor and the normal sample.
    aim: compare with a second dictionary that have all the phenotipic information. 
    
    '''    
    ifile = open(file)
    pd=defaultdict(defaultdict)
    for l in ifile:
        fields=l.strip('\n').split(' ')
        fpath=fields[-1]
        patient=fpath.split('/')[1]
        fname=fpath.split('/')[-1]
        fn = fname.split('.')[1]
        sp = fn.split('-')[3]
        if int(sp[0])==0:
            pd[patient]['tumor']=os.path.join(rootpath,fpath.lstrip('./'))
        elif int(sp[0])==1:
            pd[patient]['benign']=os.path.join(rootpath,fpath.lstrip('./'))
    return pd
        
        
def read_sample_file(file):
    '''
    '''
    ifile = open(file)
    header = ifile.next().split(',')
    samples=defaultdict(defaultdict)
    
    for l in ifile:
        fields=l.strip('\n').split(',')
        pid=fields[0]
        record=tcgaRecord(header,fields)
        if not record.is_tumor():
            samples[pid]['benign']=record
        else:
            samples[pid]['tumor']=record
        
    return samples
    


def create_analysis_files(masterfile,tcgarecords,bamfilesdict,rootdir,remoterootdir):
    print masterfile
    id=0
    for patient,muestras in bamfilesdict.iteritems():
        print patient
        
        tree= etree.parse(masterfile)
        root = tree.getroot()
        root.set("name",patient)
        r = root.find("remote_output_dir")
        r.set("path",os.path.join(remoterootdir,patient))
        samples=root.findall("sample")
        print samples
        
        spid=0
                
        for stype,filepath in muestras.iteritems():
            print stype,filepath
            record=tcgarecords[patient][stype]
            print record
            sample_name = "%s_%s"%(patient,stype)
            lib_id="SI_%d"%(id)
            lane_name = "%s_%s_%s"%("tcga",record.sample_name(),lib_id)
            id+=1
            description="%s_%s"%(record.disease(),stype)
            print sample_name, lane_name
            lanefields={'lib_id':lib_id,'description':description,
                       "aligned_bam_file":filepath,"sample_id":"%s_%s_null"%(patient,stype),
                       "center":"TCGA","read_length":"0","fragment_length_mean":"0"}
            
            sp = samples[spid]
            sp.set("name",sample_name)
            samplefields={"patient_id":patient,"category":stype}
            spid+=1
            for f,v in samplefields.iteritems():
                print f,v
                ri = sp.find(f)
                print ri.text
                ri.text=v
                print ri.text
                
            
            lane = sp.findall('lane')[0]
            lane.set("name",lane_name)
            print lane
            
            
            for f,v in lanefields.iteritems():
                print f,v
                ri = lane.find(f)
                print ri.text
                ri.text=v
                print ri.text
                
            
        sf=open(os.path.join(rootdir,"%s_analisis.xml"%(patient)),'w')
        print >> sf, etree.tostring(tree, pretty_print = True)
        sf.close()
     
    return 


def parse_xml(filename):
    
    tree= etree.parse(masterfile)
    root = tree.getroot()
    fields=['name','run_id','lib_id','description','aligned_bam_file']
    samples=root.findall("sample")
    print samples
    for sp in samples:
        lanes = sp.findall('lane')
        print lanes
        for lane in lanes:
            ri = lane.find("run_id")
            print ri.text
            ri.text="NewONe"
            print ri.text
            sys.exit(0)
            print lane.attrib
            for l in lane:
                print l
        '''
            for f in fields:
                print f
                l=lane.get(f)
                print l
                print l.text
        '''     
    
    sys.exit(0)
    
    for element in root:#.iter("*"):
        #print element
        print element.tag
        print element.text
        if element.tag == "remote_output_dir":
            element.text = "rootpath"
            print element.text
        
        if element.tag == "sample":
            pass
       

    
remoterootdir="/nobackup/med-mctp/projects/exome/tcga/"
masterfile='/exds/users/oabalbin/projects/exomes/tcga/masterfile_analysis.xml'
rootAnalysisPath="/exds/users/oabalbin/projects/exomes/tcga/analysisToRun/"
rootpath='/exds/dbgap/tcga/nov17/PRAD_decripted/22600/analysis/PRAD'
file='/exds/users/oabalbin/projects/exomes/tcga/PRAD.bam.files.decripted.txt'
sample_file = '/exds/users/oabalbin/projects/exomes/tcga/TCGA_phs000178_tranche_Q-1.txt_processed_withAll_PRAD.csv'
pd = read_bam_file_list(file,rootpath)
samples=read_sample_file(sample_file)
create_analysis_files(masterfile,samples,pd,rootAnalysisPath,remoterootdir)

#parse_xml(masterfile)